/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <cmath>
#include <iostream>
#include "vec_dot.h"
#include "parvec_axpby.h"
#include "par_pipegmres_solver.h"

namespace YHAMG
{

Par_PipeGMRES_Solver::Par_PipeGMRES_Solver(const Par_Operator& _A)
	: A(&_A),
	P(0),
	Restart(10),
	Shift(1.0),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_PipeGMRES_Solver::Par_PipeGMRES_Solver(const Par_Operator& _A, const Par_Operator& _P)
	: A(&_A),
	P(&_P),
	Restart(10),
	Shift(1.0),
	MaxIters(500),
	Tolerance(1.0e-08),
	PrintStats(0)
{
}

Par_PipeGMRES_Solver::Par_PipeGMRES_Solver(const Par_Operator& _A, int restart, double shift, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(0),
	Restart(restart),
	Shift(shift),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}

Par_PipeGMRES_Solver::Par_PipeGMRES_Solver(const Par_Operator& _A, const Par_Operator& _P, int restart, double shift, int max_Iters, double tolerance, int print_stats)
	: A(&_A),
	P(&_P),
	Restart(restart),
	Shift(shift),
	MaxIters(max_Iters),
	Tolerance(tolerance),
	PrintStats(print_stats)
{
}


#define ABS(x) ((x) > 0 ? (x) : -(x))
#define H(i, j) (H[(i) + (((j) + 1) * (j) / 2)])

static inline void GeneratePlaneRotation(double dx, double dy, double& cs, double& sn)
{
	if (dy == 0.0)
	{
		cs = 1.0;
		sn = 0.0;
	}
	else if (ABS(dy) > ABS(dx))
	{
		double temp = dx / dy;
		sn = 1.0 / sqrt(1.0 + temp * temp);
		cs = temp * sn;
	}
	else
	{
		double temp = dy / dx;
		cs = 1.0 / sqrt(1.0 + temp * temp);
		sn = temp * cs;
	}
}

static inline void ApplyPlaneRotation(double cs, double sn, double& dx, double& dy)
{
	double temp = cs * dx + sn * dy;
	dy = -sn * dx + cs * dy;
	dx = temp;
}

static inline void SolveUpperTriangular(int m, const double* H, double* y)
{
	for (int i = m - 1; i >= 0; --i)
	{
		y[i] /= H(i, i);
		for (int j = i - 1; j >= 0; --j)
			y[j] -= H(j, i) * y[i];
	}
}

#define MACHINE_EPSILON        2.2204460492503131e-16
#define APPLYDOT(x, y, result) ApplyDot(x, y, result, buf, dst, cnt)
#define BEGINREDUCTION()       BeginReduction(buf, cnt, comm, &request)
#define ENDREDUCTION()         EndReduction(buf, dst, cnt, &request, &status)

static inline void ApplyDot(const Par_Vector& x, const Par_Vector& y, double& result, double* buf, double **dst, int &cnt)
{
	buf[cnt] = Vec_Dot(x.Loc, y.Loc);
	dst[cnt++] = &result;
}

static inline void BeginReduction(double* buf, int cnt, MPI_Comm comm, MPI_Request* request)
{
	MPI_Iallreduce(MPI_IN_PLACE, buf, cnt, MPI_DOUBLE, MPI_SUM, comm, request);
}

static inline void EndReduction(const double* buf, double** dst, int &cnt, MPI_Request* request, MPI_Status* status)
{
	MPI_Wait(request, status);
	for (int i = 0; i < cnt; ++i)
		*dst[i] = buf[i];
	cnt = 0;
}

void Par_PipeGMRES_Solver::operator()(const Par_Vector& b, const Par_Vector& x, int& iter, double& relres) const
{
	MPI_Comm comm = A->comm;

	int rank;
	MPI_Comm_rank(comm, &rank);

	int m = Restart;

	MPI_Request request;
	MPI_Status status;
	double* buf = new double[m + 1];
	double** dst = new double*[m + 1];
	int cnt = 0;

	int n = A->OutSize();

	double normb, res;
	double* cs = new double[m];
	double* sn = new double[m];
	double* y = new double[m + 1];
	double* H = new double[(m + 1) * m / 2 + 1];
	Par_Vector w(comm), q(comm);
	Par_Vector* p = new Par_Vector[m];
	Par_Vector* z = new Par_Vector[m];

	w.Resize(n);
	q.Resize(n);
	for (int k = 0; k < m; ++k)
	{
		p[k].comm = comm;
		z[k].comm = comm;
		p[k].Resize(n);
		z[k].Resize(n);
	}

	if (PrintStats && rank == 0)
	{
		std::cout << "Iter\tResidual\n";
		std::cout << "----------------------------------------------------------------\n";
	}

	iter = 0;

	APPLYDOT(b, b, normb);

	BEGINREDUCTION();

	while (1)
	{
		A->Apply(x, p[0]);
		ParVec_AXPBY(-1.0, p[0], 1.0, b, p[0]);

		if (iter == 0)
		{
			ENDREDUCTION();
			normb = sqrt(normb);
		}

		APPLYDOT(p[0], p[0], res);

		BEGINREDUCTION();

		if (P)
			P->Apply(p[0], q);
		else
			q.Copy(p[0]);

		A->Apply(q, z[0]);
		z[0].AddScaled(-Shift, p[0]);

		ENDREDUCTION();

		res = sqrt(res);

		if (iter == 0 && PrintStats && rank == 0)
			std::cout << iter << "\t" << res << "\n";

		if (res / normb <= Tolerance || iter == MaxIters) break;

		p[0].Scale(1.0 / res);
		z[0].Scale(1.0 / res);

		for (int i = 0; i <= m; ++i)
			y[i] = 0.0;

		y[0] = res;

		int k = 0;

		while (1)
		{
			for (int i = 0; i <= k; ++i)
				APPLYDOT(p[i], z[k], H(i, k));

			APPLYDOT(z[k], z[k], H(k + 1, k));

			BEGINREDUCTION();

			if (k < m - 1)
			{
				if (P)
					P->Apply(z[k], q);
				else
					q.Copy(z[k]);
				A->Apply(q, z[k + 1]);
			}

			ENDREDUCTION();

			for (int i = 0; i <= k; ++i)
				H(k + 1, k) -= H(i, k) * H(i, k);

			if (H(k + 1, k) < 0 && H(k + 1, k) >= -MACHINE_EPSILON) H(k + 1, k) = 0;
			if (H(k + 1, k) < 0) break;

			H(k + 1, k) = sqrt(H(k + 1, k));

			if (k < m - 1)
			{
				p[k + 1].Copy(z[k]);
				int i = 0;
				for ( ; i < k; i += 2)
					p[k + 1].Add2Scaled(-H(i, k), p[i], -H(i + 1, k), p[i + 1]);
				if (i == k) p[k + 1].AddScaled(-H(i, k), p[i]);
				p[k + 1].Scale(1.0 / H(k + 1, k));
			}

			H(k, k) += Shift;

			if (k < m - 1)
			{
				int i = 0;
				for ( ; i < k; i += 2)
					z[k + 1].Add2Scaled(-H(i, k), z[i], -H(i + 1, k), z[i + 1]);
				if (i == k) z[k + 1].AddScaled(-H(i, k), z[i]);
				z[k + 1].Scale(1.0 / H(k + 1, k));
			}
			
			for (int i = 0; i < k; ++i)
				ApplyPlaneRotation(cs[i], sn[i], H(i, k), H(i + 1, k));

			GeneratePlaneRotation(H(k, k), H(k + 1, k), cs[k], sn[k]);
			ApplyPlaneRotation(cs[k], sn[k], H(k, k), H(k + 1, k));
			ApplyPlaneRotation(cs[k], sn[k], y[k], y[k + 1]);

			res = ABS(y[k + 1]);

			++k;
			++iter;

			if (rank == 0 && PrintStats)
				std::cout << iter << "\t" << res << "\n";

			if (k == m || res / normb <= Tolerance || iter == MaxIters) break;
		}

		SolveUpperTriangular(k, H, y);

		w.Fill(0.0);

		int i = 0;
		for ( ; i + 1 < k; i += 2)
			w.Add2Scaled(y[i], p[i], y[i + 1], p[i + 1]);
		if (i + 1 == k) w.AddScaled(y[i], p[i]);

		if (P)
			P->Apply(w, q);
		else
			q.Copy(w);

		x.AddScaled(1.0, q);

		if (res / normb <= Tolerance || iter == MaxIters) break;
	}

	relres = res / normb;

	if (PrintStats && rank == 0)
	{
		std::cout << "----------------------------------------------------------------\n";
		std::cout << "Iterations: " << iter << "\n";
		std::cout << "Final Relative Residual: " << relres << "\n";
	}

	delete[] p;
	delete[] z;
	delete[] y;
	delete[] cs;
	delete[] sn;
	delete[] H;
	delete[] buf;
	delete[] dst;
}

}